package union_set


func NewUnionSet(values []interface{}) UnionSet {
	nodesMap := map[interface{}]Node{}
	for i:=0; i<len(values); i++ {
		nodesMap[values[i]] = Node{value: values[i]}
	}
	parentMap := map[Node]Node{}
	setSize := map[Node]int{}
	for _, v := range nodesMap{
		parentMap[v] = v
		setSize[v] = 1
	}
	return UnionSet{nodesMap, parentMap, setSize}
}

type Node struct {
	value interface{}
}

type UnionSet struct {
	nodesMap map[interface{}]Node
	parentMap map[Node]Node
	SetSize map[Node]int
}

func (s *UnionSet) Union(a, b interface{}) {
	if !s.IsContain(a) || !s.IsContain(b) {
		panic(any("input invalid !!!, value not init in nodes Map"))
	}
	if !s.IsSameSet(a, b) {
		aHead := s.FindFather(s.nodesMap[a])
		bHead := s.FindFather(s.nodesMap[b])
		if s.SetSize[aHead] >  s.SetSize[bHead] {
			s.parentMap[bHead] = aHead
			s.SetSize[aHead] = s.SetSize[aHead] + s.SetSize[bHead]
			delete(s.SetSize, bHead)
		} else {
			s.parentMap[aHead] = bHead
			s.SetSize[bHead] = s.SetSize[aHead] + s.SetSize[bHead]
			delete(s.SetSize, aHead)
		}
	}
}

func (s *UnionSet) IsSameSet(a, b interface{}) bool {
	if s.IsContain(a) && s.IsContain(b) {
		return s.FindFather(s.nodesMap[a]) == s.FindFather(s.nodesMap[b])
	}
	return false
}

func (s *UnionSet) IsContain(a interface{}) bool {
	_, ok := s.nodesMap[a]
	return ok
}

func (s *UnionSet) FindFather(a Node) Node {
	// 收集找 parent 节点的路径，方便将路径扁平化
	path := []Node{}
	// 最顶层的节点的父节点，一定指向自己，只要 节点的父节点是 自己，就说明找到了父节点，直接返回即可
	for a != s.parentMap[a] {
		a = s.parentMap[a]
		path =append(path, a)
	}
	// 扁平化，直接将路径上的点直接指向parent
	for _, n := range path {
		s.parentMap[n] = a
	}
	return a
}